# coding=UTF-8

"""
Sparse Bundle Adjustment

"""


import numpy as np
import sba



def GenerateSBAdata(pts3D,tracks,cams,K):
    """
    功能：生成SBA模块所需格式的数据
    输入
        pts3D:  3D点列表，每个元素为[X,Y,Z]
        tracks: 对应3D点的PointTrack对象列表,每个PointTrack对应到
        cams:   相机列表，每个元素为[idx,np.array([.orient.]),np.array([[x],[y],[z]])]
        K:      相机内参
    输出：

    """
    """将输入的信息转换为SBA所需的参数 """
    ##数据点的转换
    NumOfCam=len(cams)      #相机的数量

    ##生成3D点以及对应视图2D图像点的坐标阵列
    N=len(pts3D)
    sbaPts=np.zeros([N,3+NumOfCam*2])
    for i in range(N):
        sbaPtRow=[]         ##SBA的Points类所需的阵列的每一行
        sbaPtRow.extend(pts3D[i])       ##填写3D点坐标
        track=tracks[i]
        ##按相机依次填写图像点坐标
        for x in cams:
            idx=x[0]
            if idx not in track.idx_list:   ##没有点轨迹上没有对应相机，则u,v设为nan
                u,v=(np.nan,np.nan)
            else:
                j=track.idx_list.index(idx)
                u,v=track.pos_list[j]       #取出对应相机的图像点坐标
            sbaPtRow.extend([u,v])
        ##将此3D点填入阵列
        sbaPts[i]=sbaPtRow

    ##生成相机的参数阵列，17个参数（NumOfCam*17的阵列）
    sbaCams=np.zeros([NumOfCam,17])
    for i in range(NumOfCam):
        x=cams[i]
        camsPara=[]
        fx=K[0,0];
        cx=K[2,0]; cy=K[2,1]    ##中心点
        AR=K[0,0]/K[1,1]        #aspect ratio
        skew=K[0,1]  ##x焦距，skew coeff
        quat=sba.quaternions.quatFromRotationMatrix(x[1])       ##将R矩阵转换为四元数
        quat=quat.asVector()
        pos=x[2].reshape(3)     ##位置，将列向量变为列表
        camsPara.extend([fx,cx,cy,AR,skew])     ##填入内参
        camsPara.extend([0,0,0,0,0])    ##因为已经处理了畸变，所以这里都为0
        camsPara.extend(quat)
        camsPara.extend(pos)
        ##将此相机参数填入阵列
        sbaCams[i]=camsPara
    cameras=sba.Cameras.fromDylan(sbaCams)
    points=sba.Points.fromDylan(sbaPts)

    return cameras,points


pts=sba.Points.fromTxt('/home/wg/7pts.txt',7)
cams=sba.Cameras.fromTxt('/home/wg/7camsvarK.txt')
newcams, newpts, info = sba.SparseBundleAdjust(cams,pts)
print('OK')
#print(pts)